import os
import unittest
import torchaudio
from lid.lid_proccesser import WenetEncoder
from lid.lid_proccesser import MyClassifier


class TestLidProgresser(unittest.TestCase):
    def setUp(self) -> None:
        model_name = "langid-ecapa-encoder-512-epoch-5"
        lid_model_dir = os.path.join(os.path.dirname(__file__), f'models/{model_name}')
        encoder_path = os.path.join(os.path.dirname(__file__), 'models/encoder.onnx')

        self.encoder = WenetEncoder(encoder_path=encoder_path)
        self.classifier = MyClassifier.from_hparams(
            source=lid_model_dir, savedir=lid_model_dir
        )
        self.classifier.set_encoder(self.encoder)
        data_dir = os.path.join(os.path.dirname(__file__), "../data")
        self.test_langs = ["en", "zh-ca", "zh-sh", "zh-si"]
        self.test_wavs = [
            os.path.join(data_dir, f"youtube_{lang}.wav") for lang in self.test_langs
        ]

    def test_compute_batch_encoder_output(self):
        waveform, sr = torchaudio.load(self.test_wavs[0])
        assert sr == self.classifier.sr

        waveform = waveform[:, 30 * 1600]
        encoder_output = self.encoder.compute_batch_encoder_output(waveform)
        self.assertEqual(len(encoder_output.shape), 3)
        self.assertEqual(encoder_output.shape[2], 512)


    def test_detect_lid(self):
        for ref_lang, wav in zip(self.test_langs, self.test_wavs):
            waveform, sr = torchaudio.sr
            hyp_lang = self.classifier.detect_sentenc_lid(waveform)
            self.assertEqual(ref_lang, hyp_lang)

if __name__ == "__main__":
    unittest.main()